# coding = gb2312
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec
import os
import time

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.size'] = 14


def load_local_image(image_path, max_size=400):
    """
    参数:
        image_path: 本地图像文件路径
        max_size: 调整图像的最大尺寸（保持宽高比）
    """
    # 使用matplotlib加载图像（返回NumPy数组）
    image = plt.imread(image_path)

    # 确保图像是uint8类型
    if image.dtype == np.float32:
        image = (image * 255).astype(np.uint8)

    # 如果图像有透明度通道，去除它
    if image.shape[2] == 4:
        image = image[:, :, :3]

    # 如果图像太大，调整大小
    height, width = image.shape[:2]
    if max(height, width) > max_size:
        scale = max_size / max(height, width)
        new_width = int(width * scale)
        new_height = int(height * scale)
        image = resize_image(image, new_width, new_height)

    # 转换为浮点数并归一化到[0,1]范围
    return image.astype(np.float32) / 255.0


def resize_image(image, new_width, new_height):
    """
    参数:
        image: 输入图像数组
        new_width: 新宽度
        new_height: 新高度
    """
    # 获取原始尺寸
    height, width, channels = image.shape

    # 创建新图像的网格
    x_indices = (np.arange(new_width) * (width / new_width)).astype(np.int32)
    y_indices = (np.arange(new_height) * (height / new_height)).astype(np.int32)

    # 使用高级索引获取像素值
    resized = image[y_indices[:, None], x_indices]

    return resized


def svd_compress_channel(channel, compression_ratio):
    """
    参数:
        channel: 单通道图像数据
        compression_ratio: 压缩率 (0-1)
    """
    # 执行SVD分解
    u, s, vt = np.linalg.svd(channel, full_matrices=False)

    # 计算需要保留的奇异值数量
    k = max(1, int(len(s) * compression_ratio))

    # 重构压缩后的通道
    compressed_channel = u[:, :k] @ np.diag(s[:k]) @ vt[:k, :]

    # 计算压缩比和存储空间节省
    original_size = channel.shape[0] * channel.shape[1]
    compressed_size = u.shape[0] * k + k + k * vt.shape[1]
    compression_ratio_actual = compressed_size / original_size
    space_saving = (1 - compression_ratio_actual) * 100

    return np.clip(compressed_channel, 0, 1), k, space_saving


def svd_compress_image(image, compression_ratio):
    """
    参数:
        image: 输入彩色图像 (H x W x 3)
        compression_ratio: 压缩率 (0-1)
    """
    compressed_channels = []
    k_values = []
    space_savings = []

    # 对每个颜色通道分别压缩
    for i in range(3):
        compressed_channel, k, space_saving = svd_compress_channel(image[:, :, i], compression_ratio)
        compressed_channels.append(compressed_channel)
        k_values.append(k)
        space_savings.append(space_saving)

    # 平均存储空间节省
    avg_space_saving = np.mean(space_savings)

    # 合并通道
    compressed_image = np.stack(compressed_channels, axis=2)

    return compressed_image, k_values, avg_space_saving


def plot_svd_effects(image, compression_ratios, output_path="svd_compression_results.png"):
    """
    参数:
        image: 原始图像
        compression_ratios: 压缩率列表
        output_path: 结果图像保存路径
    """
    # 创建背景图
    fig = plt.figure(figsize=(24, 28))
    fig.suptitle('SVD奇异值图像压缩', fontsize=26, fontweight='bold', y=1.00)

    # 使用嵌套的GridSpec创建更灵活的布局
    outer_gs = GridSpec(2, 1, height_ratios=[2.5, 1], hspace=0.3, top=0.96)  # 增加垂直间距

    # 上部区域：图像和说明
    top_gs = GridSpecFromSubplotSpec(2, 3, subplot_spec=outer_gs[0],
                                     width_ratios=[1, 1, 1.2], height_ratios=[1, 1],
                                     wspace=0.3, hspace=0.2)  # 增加子图间距

    # 下部区域：奇异值分布图
    bottom_gs = GridSpecFromSubplotSpec(1, 1, subplot_spec=outer_gs[1])

    # 原始图像
    ax0 = fig.add_subplot(top_gs[0, 0])
    ax0.imshow(image)
    ax0.set_title('原始图像', fontsize=20, fontweight='bold', pad=10)
    ax0.axis('off')

    # 压缩图像
    compression_titles = ["50%压缩", "25%压缩", "10%压缩", "1%压缩"]
    positions = [(0, 1), (0, 2), (1, 0), (1, 1)]  # 重新安排位置

    # 创建列表存储压缩结果
    compression_results = []

    for ratio in compression_ratios:
        # 执行压缩并计时
        start_time = time.perf_counter()  # 使用标准time模块
        compressed_img, k_values, space_saving = svd_compress_image(image, ratio)
        compression_time = time.perf_counter() - start_time
        compression_results.append((compressed_img, k_values, space_saving, compression_time))

    # 显示压缩后的图像
    for i, (row, col) in enumerate(positions):
        compressed_img, k_values, space_saving, compression_time = compression_results[i]
        ax = fig.add_subplot(top_gs[row, col])
        ax.imshow(compressed_img)

        # 显示压缩信息
        title = f'{compression_titles[i]}\n奇异值: R{k_values[0]}/G{k_values[1]}/B{k_values[2]}'
        ax.set_title(title, fontsize=16, pad=8)

        # 在图像下方添加额外信息
        info_text = f'空间节省: {space_saving:.1f}%\n耗时: {compression_time:.2f}s'
        ax.text(0.5, -0.005, info_text, transform=ax.transAxes,  # 降低文本位置
                ha='center', va='top', fontsize=14)
        ax.axis('off')

    # 添加SVD解释说明
    ax_text = fig.add_subplot(top_gs[1, 2])
    explanation = (
        "SVD奇异值分解图像压缩原理\n\n"
        "1. 数学基础:\n"
        "    矩阵分解: A = UV（转置）\n"
        "    U: m×m 正交矩阵\n"
        "    Σ: 对角矩阵(奇异值)\n"
        "    V（转置）: 正交矩阵的转置\n\n"
        "2. 压缩原理:\n"
        "    保留前k个最大的奇异值\n"
        "    重构近似矩阵: A = UV（转置）\n\n"
        "3. 存储优化:\n"
        "    原始存储: m×n\n"
        "    压缩后存储: k×(m+n+1)\n\n"
        "4. 压缩效果:\n"
        "    50%: 几乎无损\n"
        "    25%: 轻微细节损失\n"
        "    10%: 明显细节损失\n"
        "    1%: 高度抽象化"
    )
    ax_text.text(0.05, 0.5, explanation, fontsize=16, ha='left', va='center',
                 bbox=dict(boxstyle="round,pad=0.5", facecolor='#f0f8ff', alpha=0.8))
    ax_text.axis('off')

    # 添加奇异值分布图 - 使用底部整个区域
    ax_svd = fig.add_subplot(bottom_gs[0])
    colors = ['#ff0000', '#00cc00', '#0000ff']  # 红, 绿, 蓝
    labels = ['红色通道', '绿色通道', '蓝色通道']
    line_styles = ['-', '--', '-.']
    markers = ['o', 's', '^']  # 添加标记点

    # 计算并绘制累积能量
    max_singular_values = 0
    singular_values = []
    for i in range(3):
        u, s, vt = np.linalg.svd(image[:, :, i], full_matrices=False)
        singular_values.append(s)
        max_singular_values = max(max_singular_values, len(s))
        cumulative_energy = np.cumsum(s) / np.sum(s)

        # 稀疏采样以减少曲线密度
        step = max(1, len(s) // 100)
        indices = np.arange(0, len(s), step)
        if indices[-1] != len(s) - 1:
            indices = np.append(indices, len(s) - 1)

        ax_svd.plot(indices, cumulative_energy[indices],
                    label=labels[i], color=colors[i],
                    lw=2, linestyle=line_styles[i],
                    marker=markers[i], markersize=4, markevery=20)

    # 标记压缩率位置 - 使用垂直虚线并动态调整标签位置
    label_positions = [0.95, 0.85, 0.75, 0.65]  # 预定义垂直位置

    for i, ratio in enumerate(compression_ratios):
        xpos = ratio * max_singular_values

        # 计算实际位置
        actual_x = min(xpos, max_singular_values - 1)
        ax_svd.axvline(x=actual_x, color='gray', linestyle=':', alpha=0.7, lw=1.5)

        # 在曲线顶部添加标签
        ax_svd.text(actual_x, label_positions[i], f'{int(ratio * 100)}%压缩', fontsize=14,
                    ha='center', va='center', rotation=0,
                    bbox=dict(boxstyle="round,pad=0.3", facecolor='white', alpha=0.9))

    ax_svd.set_title('奇异值累积能量分布', fontsize=20, fontweight='bold', pad=15)
    ax_svd.set_xlabel('保留的奇异值数量', fontsize=16, labelpad=10)
    ax_svd.set_ylabel('累积能量百分比', fontsize=16, labelpad=10)
    ax_svd.grid(True, linestyle='--', alpha=0.6)
    ax_svd.legend(fontsize=14, loc='lower right', framealpha=0.9)
    ax_svd.set_ylim(0, 1.02)
    ax_svd.set_xlim(0, max_singular_values)
    ax_svd.tick_params(axis='both', labelsize=12)

    # 保存高质量结果图像
    plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')
    print(f"结果已保存至: {os.path.abspath(output_path)}")
    plt.show()


def main():
    image_path = "D:/python/OIP-C.webp"

    # 检查文件是否存在
    if not os.path.exists(image_path):
        print(f"错误: 图像文件不存在 - {image_path}")
        return

    print(f"正在处理图像: {os.path.abspath(image_path)}")

    # 加载图像
    image = load_local_image(image_path, max_size=400)
    print(f"图像尺寸: {image.shape[1]}x{image.shape[0]}")

    # 设置压缩率
    compression_ratios = [0.5, 0.25, 0.1, 0.01]

    # 可视化压缩效果
    plot_svd_effects(image, compression_ratios)


if __name__ == "__main__":
    main()